import re
import pandas as pd
from nltk.corpus import stopwords
from sklearn.model_selection import train_test_split
from sparkai.llm.llm import ChatSparkLLM
from sparkai.core.messages import ChatMessage

# 星火认知大模型配置
SPARKAI_URL = "wss://spark-api.xf-yun.com/v4.0/chat"
SPARKAI_APP_ID = "ad70c071"
SPARKAI_API_SECRET = "YWNjZjNlY2VhOGQ4ODhhYTVmZTMxN2Y3"
SPARKAI_API_KEY = "07bc2274aa170e50483ba19ae13030f1"
SPARKAI_DOMAIN = "4.0Ultra"

stop_words = set(stopwords.words("english"))


class TopMiningMax:
    def __init__(self, filename):
        # 读取数据
        data = pd.read_csv("../InputData/" + filename + ".tsv", sep="\t")
        # 删除缺失值
        data.dropna(
            subset=["benefitsReview", "sideEffectsReview", "commentsReview"],
            inplace=True,
        )
        # 应用预处理
        data["benefitsReview_cleaned"] = data["benefitsReview"].apply(
            self.preprocess_text
        )
        data["sideEffectsReview_cleaned"] = data["sideEffectsReview"].apply(
            self.preprocess_text
        )
        data["commentsReview_cleaned"] = data["commentsReview"].apply(
            self.preprocess_text
        )
        # 合并文本字段
        data["combinedText"] = (
            data["benefitsReview_cleaned"]
            + " "
            + data["sideEffectsReview_cleaned"]
            + " "
            + data["commentsReview_cleaned"]
        )
        # 将文本和药品名称分开
        self.X = data["combinedText"]
        self.y = data["urlDrugName"]
        # 创建药品名称列表
        drug_names = data["urlDrugName"]
        self.drug_list = list(set(drug_names))
        # 划分训练和测试集
        self.X_train, self.X_test, self.y_train, self.y_test = train_test_split(
            self.X, self.y, test_size=0.2, random_state=42
        )

    def extract_drug_name_from_comment(self, comment, drug_list):
        """
        从药物评论中提取药物名称。

        参数:
        comment (str): 药物评论文本。
        drug_list (list): 包含药物名称的列表。

        返回:
        str: 提取的药物名称。
        """
        # 初始化星火大模型
        spark = ChatSparkLLM(
            spark_api_url=SPARKAI_URL,
            spark_app_id=SPARKAI_APP_ID,
            spark_api_key=SPARKAI_API_KEY,
            spark_api_secret=SPARKAI_API_SECRET,
            spark_llm_domain=SPARKAI_DOMAIN,
            streaming=False,
        )
        # 准备输入消息，包括提示信息
        prompt_message = "请根据我提供药物名称列表，结合给出的患者有关药物的评论，选出最符合的药品名称，只回复药物名称(如果推理不出，随机选择一个)."
        drug_options = ", ".join(drug_list)  # 将药物名称列表转换为字符串，用逗号分隔
        messages = [
            ChatMessage(
                role="user", content=prompt_message + " " + drug_options + " " + comment
            )
        ]

        # 调用星火大模型生成回复
        response = spark.generate([messages])

        # 提取药物名称
        try:
            # 检查 generations 属性
            if hasattr(response, "generations") and response.generations:
                # 获取第一个生成的结果
                first_generation = response.generations[0]
                # 检查 first_generation 的内容
                if first_generation:
                    drug_text = first_generation[
                        0
                    ].text.strip()  # 只有一个生成，取第一个
                    # 将药品名称列表转换成正则表达式模式
                    # 使用 re.escape 确保药品名称中的特殊字符不会影响正则表达式
                    pattern = "|".join([re.escape(drug) for drug in drug_list])

                    # 在文本中查找所有匹配的药品名称
                    drug_name = re.findall(pattern, drug_text, flags=re.IGNORECASE)
                else:
                    drug_name = "None"
            else:
                drug_name = "None"
        except Exception as e:
            print(f"Error extracting drug name: {e}")
            drug_name = "None"
        # 将列表去重并转换为字符串
        results = list(set(drug_name))
        # 提取药名
        if len(results) == 0:
            result = "None"
        else:
            result = results[0]
        return result.lower()

    # 文本清理函数
    def preprocess_text(self, text):
        if pd.isna(text):  # 检查是否为 NaN
            return ""  # 返回空值作为默认值
        # 小写化
        text = text.lower()
        # 去除特殊字符和数字
        text = re.sub(r"[^a-zA-Z\s]", "", text)
        # 去掉停用词
        text = " ".join(word for word in text.split() if word not in stop_words)
        return text

    def extract_features(self):
        # 示例：测试集的前10条数据
        X_test_sample = self.X_test[:10].tolist()  # 转为Python列表
        true_drugs_sample = self.y_test[:10].values  # 获取真实药物名称

        # 逐条预测药物名称
        api_predictions = [
            self.extract_drug_name_from_comment(text, self.drug_list)
            for text in X_test_sample
        ]

        # 对比预测结果
        comparison = pd.DataFrame(
            {"TrueDrugName": true_drugs_sample, "PredictedDrugName": api_predictions}
        )

        print(comparison)

        # 计算 API 准确率
        api_accuracy = sum(
            comparison["TrueDrugName"] == comparison["PredictedDrugName"]
        ) / len(comparison)
        print("API Prediction Accuracy:", api_accuracy)


if __name__ == "__main__":
    Max = TopMiningMax("drugLibTest_raw")
    Max.extract_features()
